#!/usr/bin/env python2.5_cfd
# Original script received from "Hans-Wilfried Mindt" <hmi@esigmbh.de> on Jul 29, 2009
# Modified by "Brett Tully" <brett.tully@eng.ox.ac.uk> on Sep 2, 2009 to suit fbg-cluster.eng.ox.ac.uk
import os, sys, glob, re, pwd, getpass, time

FBG_mem_per_node = 8
FBG_proc_per_node = 8

maxMemPerNode = FBG_mem_per_node

print '''
=====================================================================
|                  create_sge_script.py                             |
|                                                                   |
| This program generates a SGE script for the ESI-CFD fbg-cluster   |
| in its current dedicated-queue configuration.                     |
|                                                                   |
| Modified by Brett Tully to suit fbg-cluster.eng.ox.ac.uk          |
=====================================================================
'''
print ">>> How much memory (in GB) do you need for this job? [Don't know]"
memory_required = 0.0
try:
  memory_required = float(raw_input())
except ValueError:
  print ">>> You didn't enter a valid number.  Assuming memory isn't a concern."

nodes_required = int(memory_required/maxMemPerNode)
if memory_required%maxMemPerNode != 0: nodes_required += 1
if nodes_required == 0: nodes_required=1

print 
print ">>> How many processors do you want for this job? "
print "    [Just hit return to calculate the best number.]"
procs_required = FBG_proc_per_node * nodes_required
try:
  procs_required = int(raw_input())
except ValueError: pass

if procs_required < nodes_required:
  print "Sorry, there is no way to do what you've asked.  You only asked"
  print "for %d processors, but you need at least %d nodes to meet your" \
        % (procs_required,nodes_required)
  print "memory requirements.  Try again."
  sys.exit(1)
elif procs_required > FBG_proc_per_node*nodes_required:
  nodes_required = procs_required / FBG_proc_per_node + min(procs_required % FBG_proc_per_node,1)

# algorithm: 9 procs will prefer to go as 8:1 rather than 7:2
memory_per_proc = memory_required / procs_required
max_procs_per_node = FBG_proc_per_node
memory_per_node = memory_per_proc * max_procs_per_node
while memory_per_node > FBG_mem_per_node:
  max_procs_per_node = max_procs_per_node - 1
  memory_per_node = memory_per_proc * max_procs_per_node 

ppn = nodes_required * [1] + [0]
for node in range(len(ppn)):
  if sum(ppn) >= procs_required: break
  ppn[node] = ppn[node] + min( procs_required-sum(ppn), max_procs_per_node-1 )
  
SGEDirectives = []
SGEDirectives.append("#$ -pe mpich %d" % (nodes_required))
SGEDirectives.append("#$ -cwd")
SGEDirectives.append("#$ -V")
SGEDirectives.append("#$ -j y")
SGEDirectives.append("#$ -o logfile.$JOB_ID")

# Collect his inputs regarding the SGE script
candidateModels = glob.glob('*.DTF')
if not candidateModels:
  print 'There do not seem to be any DTF files in the current directory.'
  defaultModel = 'No default'
else:
  defaultModel = candidateModels[0]

answer = raw_input('>>> What model do you want to run? [%s]\n' % defaultModel )
if not answer:
  modelName = defaultModel
else:
  modelName = answer

splitName = os.path.splitext(modelName)
if ((splitName[1] == '.DTF') or (splitName[1] == '.dtf')):
  dtfFile = modelName
  modelName = splitName[0]
else:
  if os.path.isfile( modelName + '.dtf' ):
    dtfFile = modelName + '.dtf'
  elif os.path.isfile( modelName + '.DTF' ):
    dtfFile = modelName + '.DTF'
  else:
    print 
    print 'Specified file did not exist in current directory.'
    print 'I looked for %s and %s and didn\'t find either one!' % (modelName + '.DTF',
                                                                   modelName + '.dtf')
    print 
    sys.exit(1)

sendEmail = False
#whenEmail = 'e'
#while 1:
#  print '>>> Want SGE to send you an e-mail when the job finishes?'
#  print '    [option 1 - Yes, only when the job finishes.]'
#  print '    [option 2 - Yes, at job start and at job finish.]'
#  print '    [option 3 - No thanks.]'
#  answer = raw_input('>>> Enter your choice: [3]\n')
#  if '1' in answer:
#    sendEmail = True
#    whenEmail = 'e'
#    break
#  elif '2' in answer:
#    sendEmail = True
#    whenEmail = 'be'
#    break
#  elif not answer or '3' in answer:
#    sendEmail = False
#    break
#  else:
#    print 'Sorry. You have to answer 1, 2 or 3.  Those are the only options.  Try again.'

if sendEmail:
  while 1:
    print '>>> What e-mail address?'
#BJT     print '>>>     [option 1 - Send to %s]' % emailAddress2
#BJT     print '>>>     [option 2 - Send to %s@hugo (look at your .forward file)]' % emailAddress1
#BJT     print '>>>     [option 3 - other]'
    answer = raw_input('>>> Enter your choice: <No Default>\n')
    if not answer:
      print 'Sorry. You have to answer with a proper email address.  Try again.'
    else:
      emailAddress = answer
      break
  SGEDirectives.append( '#$ -m ' + whenEmail )
  SGEDirectives.append( '#$ -M %s' % emailAddress )
  
runDecompose = False
if procs_required > 1:
  answer = raw_input('>>> Skip dtf_decompose (if you select No, metis decomposition will take place)? [No]\n')
  if re.match('[y][e]*[s]*', answer, re.IGNORECASE):
    runDecompose = False
  else:
    runDecompose = True

if not runDecompose and procs_required > 1:
  try:
    dtfOutput = os.popen( 'DTF -i ' + dtfFile, 'r' )
    lines = dtfOutput.readlines()
    rc = dtfOutput.close()
    dtfCommandWorked = True
  except:
    dtfCommandWorked = False

  if dtfCommandWorked:
    nZones = -1
    for line in lines:
      if 'zone(s)' in line:
        nZones = int( line.split()[0] )
    if nZones != procs_required and nZones != -1:
      print 
      print 'You\'re going to have a problem with this case.  You said to skip'
      print 'dtf_decompose, but the number of zones in your DTF file is not the'
      print 'same as the number of processors you specified.  So cfd-ace-mpi is'
      print 'going to quit with an error.'
      print
      print 'To fix this, edit the generated *.sge file and add a \"-sim <n>\"'
      print 'where <n> is the simulation number that you would like to use.'
      print 
      answer = raw_input('>>> Go ahead anyway? [Yes]\n')
      if re.match('[n][o]*', answer, re.IGNORECASE):
        sys.exit(1)

if len(modelName) <= 11: 
  scriptName = modelName + '.sge'
else:
  scriptName = modelName[0:11] + '.sge'

answer = raw_input('>>> What do you want to call the SGE script? [%s]\n' % scriptName )
if not answer:
  pass
else:
  scriptName = answer

# Figure out type of login shell
user = getpass.getuser()
pwdEntry = pwd.getpwnam( user )
userShell = pwdEntry[6]
if 'csh' in userShell:
  csh = True; sh = False
else:
  csh = False; sh = True

# Start spitting out the script
stdout = sys.stdout
sys.stdout = open( scriptName, 'w+' )

# first piece is the SGE string
print '''\
# Automatically-generated SGE script for running CFD-ACE+ 
# Resource definitions - the following line(s) are NOT comment(s)
%s
''' % ( '\n'.join(SGEDirectives) )

# third piece is to run the job
if procs_required > 1:
  if runDecompose:
    decompOption = '-decomp -metis'
  else:
    decompOption = '-nodecomp'
  
  print '''\
WDIR=`pwd`
MODEL_NAME="%s"
HOST_FILE="$MODEL_NAME.hst"
NZONES=%d
DECOMP="%s"

# create the hostfile for CFD-ACE
cat $PE_HOSTFILE | awk '{printf("%s\\n",$1)}' > $HOST_FILE

# run CFD-ACE
cfd-ace-mpi -model $MODEL_NAME -wd $WDIR -hosts $HOST_FILE -num $NZONES $DECOMP -job -nativempi
'''%(modelName,procs_required,decompOption,"%s")

else:
  print 'CFD-ACE-SOLVER -dtf ' + dtfFile

sys.stdout.close()
sys.stdout = stdout
print '''
>>> OK.  Your SGE script %s has been created.
>>> To submit the job to the queue, enter (without quotes) "qsub %s"
>>> It is recommended that this is done from a temporary folder on the the /tmp/$USER partition
''' %(scriptName,scriptName)
print

# Check for super-obvious issues that can be caught before job submission...
sys.stdout.write('>>> Performing a few quick checks.')
sys.stdout.flush()
for i in range(10):
  time.sleep(0.1)
  sys.stdout.write('.')
  sys.stdout.flush()

if os.path.exists( modelName + '.RUN' ):
  print 
  print 'Hey, things don\'t look right!  There\'s a file called %s' % (modelName + '.RUN')
  print 'in your directory, which means there\'s a job already running.  It\'s highly'
  print 'likely your job is going to die immediately after starting.'
  print
  print '>>> First stop the running job or delete the .RUN file before you submit the job.  Bye.'
  print
  sys.exit(1)

try:
  dtfOutput = os.popen( 'DTF -vd ' + dtfFile, 'r' )
  lines = dtfOutput.readlines()
  rc = dtfOutput.close()
  dtfCommandWorked = True
except:
  dtfCommandWorked = False

warningsFound = False
if dtfCommandWorked:
  dtfVariables = { 'NEWRUN':'0', 'USER_SHARED_LIBRARY':None, 'restart_file':None }
  for line in lines:
    for key in dtfVariables.keys():
      if key in line:
        dtfVariables[key] = ' '.join(line.split()[3:])

  if not dtfVariables['restart_file']:
    dtfVariables['restart_file'] = 'No DTF sim data named restart_file.'

  if dtfVariables['NEWRUN'] == '0' or dtfVariables['NEWRUN'] == '2':
    if not os.path.isfile(dtfVariables['restart_file']):  # check if it looks like we would need a restart file
      sys.stdout.write('\n')
      print 
      print 'WARNING! Could not find the specified restart file: ' + `dtfVariables['restart_file']`
      print 
      warningsFound = True

  userlib = dtfVariables['USER_SHARED_LIBRARY']
  if userlib:
    if 'libUserAce' not in userlib:  # check it if it looks like a non-default value
      if not os.path.isfile(userlib):
        if not warningsFound: sys.stdout.write('\n')
        print 'WARNING! Could not find the specified usersub file: ' + `userlib`
        print 
        warningsFound = True

if not warningsFound:
  sys.stdout.write('OK.\n')
else:
  sys.stdout.write('...checks finished with warnings.  Please check before submitting the job.\n')
  sys.exit(1)
